/**
 * __device__ void reduce_dim0_fp32(int in_addr_, int dim1, int dim0, int
 * out_addr_)
 *
 * @details The data type is fp32, and the input data would be reshape
 *           to (dim1, dim0).
 *
 * @brief Applies the function:
 *        The input data (dim1, dim0) would be reduced_sum into (dim1, 1).
 *
 * @param in_addr_ The starting address of the input data, must be 128Byte
 *                 aligned.
 * @param dim1 The value of dim1 in reshaped dimension (dim1, dim0).
 * @param dim0 The value of dim0 in reshaped dimension (dim1, dim0). dim0 == 32X
 * @param out_addr_ The starting address of the output data, must be 128Byte
 *                  aligned.
 * @attention The space from the end of the input to the 128Byte alignment must
 *            be readable, and the space from the end of the output to the
 *            128Byte alignment must be writable.
 *            All functions in the current file follow this rule.
 */

__device__ void reduce_matrix_to_scalar(float *in,
                                        int rows,
                                        int cols,
                                        float *out_scalar)
{
    // ------------------------------------------------------------------
    // 1. row-wise reduction : (rows, cols) -> (rows, 1)
    // ------------------------------------------------------------------
    // temporary buffer holding one float per row
    __valigned__ float row_sums[32];
    reduce_dim0_fp32(reinterpret_cast<uintptr_t>(in), rows, cols,
                     reinterpret_cast<uintptr_t>(row_sums));

    // ------------------------------------------------------------------
    // 2. reduce the row sums : (rows, 1) -> (1, 1)
    // ------------------------------------------------------------------
    reduce_dim0_fp32(reinterpret_cast<uintptr_t>(row_sums), 1, rows,
                     reinterpret_cast<uintptr_t>(out_scalar));
}

// ------------------------------------------------------------------
// quick test kernel
// ------------------------------------------------------------------
__global__ void test_reduce(float *d_mat, int rows, int cols, float *d_result)
{
    reduce_matrix_to_scalar(d_mat, rows, cols, d_result);
}
